import numpy as np
import matplotlib.pyplot as plt
from mnist.multi_layer_net import MultiLayerNet
from mnist_data import load_mnist
from optimizer import SGD
from smooth_curve import *
(x_train, t_train), (x_test, t_test) = load_mnist(flatten=True, normalize=True,one_hot_label=True)
train_size=x_train.shape[0]
batch_size = 100
max_iterations=2000

#(1)进行实验的设置
#这个weight_init_types好重要啊，简直是我们三种网络的列表吗不是
#发现我们现在做项目，字典是最重要的，能储存东西很清楚
weight_init_types={'std=0.01':0.01,'Xavier':'sigmoid','He':'relu'}#宝子我就这样设置三种设置网络参数的方法
optimizer=SGD(lr=0.01)
networks={}#因为我们建立了三个网络，所以networks也是字典
train_loss={}
train_accuracy={}
#建立了这三种网络，并且保存 loss，可以选择 sigmoid 或者 rule 作为激活函数
for key,weight_init_type in weight_init_types.items():#key来表示网络名，weight_init_types传给构造函数让他来初始化数据
    networks[key]=MultiLayerNet(input_size=784,hidden_size_list=[100,100,100,100],output_size=10,activation='relu',weight_init_std=weight_init_type)
    train_loss[key]=[]
#对这三个网络都迭代了max_iterations次
#(2)开始训练
for i in range(max_iterations):
    #选出每一批的数据,mask是掩码的意思
    batch_mask=np.random.choice(train_size,batch_size)
    x_batch=x_train[batch_mask]
    t_batch=t_train[batch_mask]

    #对每一个计算出梯度并且更新梯度，同时记录一下loss的值
    for key in weight_init_types.keys():
        grads=networks[key].gradient(x_batch,t_batch)#计算梯度
        optimizer.update(networks[key].params,grads)#更新梯度
        loss=networks[key].loss(x_batch,t_batch)#计算损失函数，好记忆画图
        train_loss[key].append(loss)
    if i%100==0:
        print("=========="+'iteration:'+str(i)+'==========')
        for key in weight_init_types.keys():
            loss=networks[key].loss(x_batch,t_batch)
            print(key+":"+str(loss))

#（3）绘制图形
markers={'std=0.01':'o','Xavier':'s','He':'D'}#代表画的图的三个形状
x=np.arange(max_iterations)#迭代次数
#x：x轴的值，即为迭代次数；smooth_curve(train_loss[key]):y 轴的值，即为每一轮损失函数的值
#marker：数据点的标记样式; markevery:控制标记的密度，每隔100个数据点才画一个标记，不至于糊成一团：
#label:设置标签
for key in weight_init_types.keys():
    plt.plot(x,smooth_curve(train_loss[key]),marker=markers[key],markevery=100,label=key)
plt.xlabel('iterations')
plt.ylabel('loss')
plt.ylim(0,2.5)
plt.legend()#读取label，生成图例
plt.show()